#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
MCP训练工具核心实现
基于 Model Context Protocol 标准实现
"""

import sys
import json
import asyncio
import logging
import os
import subprocess
from typing import Dict, Any, List, Optional

# 配置日志
def setup_logging():
    """设置日志配置"""
    log_formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建logger
    logger = logging.getLogger("MCP-Training-Tool")
    logger.setLevel(logging.INFO)
    
    # 清除现有的handlers
    logger.handlers.clear()
    
    # 添加控制台handler
    console_handler = logging.StreamHandler(sys.stderr)
    console_handler.setFormatter(log_formatter)
    logger.addHandler(console_handler)
    
    # 尝试添加文件handler
    try:
        # 使用项目目录下的logs目录存储日志文件
        project_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        logs_dir = os.path.join(project_dir, "logs")
        
        # 如果logs目录不存在，则创建它
        if not os.path.exists(logs_dir):
            os.makedirs(logs_dir)
            
        log_file = os.path.join(logs_dir, "mcp_training_tool.log")
        file_handler = logging.FileHandler(log_file, encoding="utf-8")
        file_handler.setFormatter(log_formatter)
        logger.addHandler(file_handler)
        logger.info(f"日志文件创建成功: {os.path.abspath(log_file)}")
    except Exception as e:
        logger.warning(f"无法创建日志文件: {e}")
    
    return logger

logger = setup_logging()


class MiniLLMTrainingMCP:
    """mini-LLM MCP训练工具实现"""
    
    def __init__(self):
        self.initialized = False
        self.client_info = {}
        self.logger = logger
        self.project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
    
    def get_tools(self) -> List[Dict[str, Any]]:
        """获取工具列表"""
        return [
            {
                "name": "run_fine_tuning",
                "description": "运行模型微调任务",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "use_lora": {
                            "type": "boolean",
                            "description": "是否使用LoRA微调",
                            "default": True
                        },
                        "max_length": {
                            "type": "integer",
                            "description": "最大序列长度",
                            "default": 256
                        },
                        "num_train_epochs": {
                            "type": "integer",
                            "description": "训练轮数",
                            "default": 4
                        },
                        "per_device_train_batch_size": {
                            "type": "integer",
                            "description": "每个设备的批次大小",
                            "default": 1
                        },
                        "learning_rate": {
                            "type": "number",
                            "description": "学习率",
                            "default": 2e-4
                        },
                        "gradient_accumulation_steps": {
                            "type": "integer",
                            "description": "梯度累积步数",
                            "default": 8
                        },
                        "use_cpu": {
                            "type": "boolean",
                            "description": "是否使用CPU训练",
                            "default": False
                        },
                        "lora_rank": {
                            "type": "integer",
                            "description": "LoRA秩参数",
                            "default": 8
                        }
                    },
                    "required": []
                }
            },
            {
                "name": "analyze_training_data",
                "description": "分析训练数据长度分布",
                "inputSchema": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        ]
    
    async def send_response(self, request_id: Any, result: Optional[Dict[str, Any]] = None):
        """发送成功响应"""
        response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": result or {}
        }
        await self._send_json(response)
    
    async def send_error(self, request_id: Any, code: int, message: str):
        """发送错误响应"""
        error = {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": code,
                "message": message
            }
        }
        await self._send_json(error)
    
    async def _send_json(self, data: Dict[str, Any]):
        """发送JSON数据"""
        try:
            json_str = json.dumps(data, ensure_ascii=False)
            print(json_str, flush=True)
        except Exception as e:
            print(f"发送消息错误: {e}", file=sys.stderr, flush=True)
    
    async def handle_initialize(self, request_id: Any, params: Dict[str, Any]):
        """处理初始化请求"""
        try:
            # 保存客户端信息
            self.client_info = params
            
            # 检查协议版本
            protocol_version = params.get("protocolVersion", "")
            if protocol_version < "2024-01-01":
                await self.send_error(request_id, -32000, "不支持的协议版本")
                return
            
            # 返回初始化响应 - 必须包含所有必需的capabilities
            await self.send_response(request_id, {
                "protocolVersion": "2024-01-01",
                "capabilities": {
                    "tools": {},
                    "logging": {},
                    "prompts": {}
                },
                "serverInfo": {
                    "name": "mini-llm-training-mcp",
                    "version": "1.0.0"
                }
            })
            
        except Exception as e:
            await self.send_error(request_id, -32603, f"初始化错误: {str(e)}")
    
    async def handle_initialized(self):
        """处理初始化完成通知"""
        self.initialized = True
        self.logger.info("客户端初始化完成")
        
        # 发送工具列表通知
        await self._send_json({
            "jsonrpc": "2.0",
            "method": "notifications/tools/list",
            "params": {
                "tools": self.get_tools()
            }
        })
    
    async def handle_tools_list(self, request_id: Any):
        """处理工具列表请求"""
        try:
            await self.send_response(request_id, {
                "tools": self.get_tools()
            })
        except Exception as e:
            await self.send_error(request_id, -32603, f"工具列表错误: {str(e)}")
    
    async def run_fine_tuning_task(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """运行微调任务"""
        try:
            # 构建命令参数
            cmd = ["python", "training/scripts/fine_tune.py"]
            
            # 添加数据路径和输出目录
            cmd.extend([
                "--data_path", "training/data/tianqi_jiyuan_training_data.jsonl",
                "--output_dir", "training/models/finetuned_model_tianqi"
            ])
            
            # 根据参数添加选项
            if params.get("use_lora", True):
                cmd.append("--use_lora")
                
            if "max_length" in params:
                cmd.extend(["--max_length", str(params["max_length"])])
                
            if "num_train_epochs" in params:
                cmd.extend(["--num_train_epochs", str(params["num_train_epochs"])])
                
            if "per_device_train_batch_size" in params:
                cmd.extend(["--per_device_train_batch_size", str(params["per_device_train_batch_size"])])
                
            if "learning_rate" in params:
                cmd.extend(["--learning_rate", str(params["learning_rate"])])
                
            if "gradient_accumulation_steps" in params:
                cmd.extend(["--gradient_accumulation_steps", str(params["gradient_accumulation_steps"])])
                
            if params.get("use_cpu", False):
                cmd.append("--use_cpu")
                
            if "lora_rank" in params:
                cmd.extend(["--lora_rank", str(params["lora_rank"])])
            
            # 设置环境变量
            env = os.environ.copy()
            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
            
            self.logger.info(f"执行训练命令: {' '.join(cmd)}")
            
            # 执行命令
            process = subprocess.Popen(
                cmd,
                cwd=self.project_root,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            # 等待进程完成
            stdout, stderr = process.communicate()
            
            if process.returncode == 0:
                return {
                    "success": True,
                    "message": "训练任务完成",
                    "stdout": stdout,
                    "stderr": stderr
                }
            else:
                return {
                    "success": False,
                    "message": f"训练任务失败，返回码: {process.returncode}",
                    "stdout": stdout,
                    "stderr": stderr
                }
                
        except Exception as e:
            return {
                "success": False,
                "message": f"执行训练任务时出错: {str(e)}",
                "stdout": "",
                "stderr": str(e)
            }
    
    async def analyze_training_data_task(self) -> Dict[str, Any]:
        """分析训练数据"""
        try:
            # 构建命令
            cmd = ["python", "analysis_tools/analyze_data_length.py"]
            
            self.logger.info(f"执行数据分析命令: {' '.join(cmd)}")
            
            # 执行命令
            process = subprocess.Popen(
                cmd,
                cwd=self.project_root,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            # 等待进程完成
            stdout, stderr = process.communicate()
            
            if process.returncode == 0:
                return {
                    "success": True,
                    "message": "数据分析完成",
                    "stdout": stdout,
                    "stderr": stderr
                }
            else:
                return {
                    "success": False,
                    "message": f"数据分析失败，返回码: {process.returncode}",
                    "stdout": stdout,
                    "stderr": stderr
                }
                
        except Exception as e:
            return {
                "success": False,
                "message": f"执行数据分析时出错: {str(e)}",
                "stdout": "",
                "stderr": str(e)
            }
    
    async def handle_tools_call(self, request_id: Any, params: Dict[str, Any]):
        """处理工具调用请求"""
        try:
            name = params.get("name")
            arguments = params.get("arguments", {})
            
            if name == "run_fine_tuning":
                # 运行微调任务
                result = await self.run_fine_tuning_task(arguments)
                
                if result["success"]:
                    await self.send_response(request_id, {
                        "content": [
                            {
                                "type": "text",
                                "text": f"## 训练任务完成\n\n```\n{result['stdout']}\n```"
                            }
                        ]
                    })
                else:
                    await self.send_error(request_id, 1002, f"训练任务失败: {result['message']}\n\nstderr:\n{result['stderr']}")
            
            elif name == "analyze_training_data":
                # 分析训练数据
                result = await self.analyze_training_data_task()
                
                if result["success"]:
                    await self.send_response(request_id, {
                        "content": [
                            {
                                "type": "text",
                                "text": f"## 训练数据长度分析\n\n```\n{result['stdout']}\n```"
                            }
                        ]
                    })
                else:
                    await self.send_error(request_id, 1003, f"数据分析失败: {result['message']}\n\nstderr:\n{result['stderr']}")
                    
            else:
                await self.send_error(request_id, 1001, f"未知工具: {name}")
                
        except Exception as e:
            await self.send_error(request_id, -32603, f"工具执行错误: {str(e)}")
    
    async def handle_ping(self, request_id: Any):
        """处理ping请求，保持连接活跃"""
        self.logger.info("收到ping请求，发送pong响应")
        await self.send_response(request_id, {"pong": True})
    
    async def handle_request(self, message: Dict[str, Any]):
        """处理客户端请求"""
        try:
            method = message.get("method")
            request_id = message.get("id")
            params = message.get("params", {})
            
            if method == "initialize":
                await self.handle_initialize(request_id, params)
            elif method == "tools/list":
                await self.handle_tools_list(request_id)
            elif method == "tools/call":
                await self.handle_tools_call(request_id, params)
            elif method == "initialized":
                await self.handle_initialized()
            elif method == "ping":
                await self.handle_ping(request_id)
            elif method == "$/cancelRequest":
                # 处理取消请求
                pass
            else:
                await self.send_error(request_id, -32601, f"未知方法: {method}")
                
        except Exception as e:
            error_id = message.get("id")
            if error_id:
                await self.send_error(error_id, -32603, f"请求处理错误: {str(e)}")
    
    async def run(self):
        """运行服务器主循环"""
        try:
            self.logger.info("开始初始化训练工具")
            self.logger.info(f"项目根目录: {self.project_root}")
            self.logger.info("训练工具初始化完成，等待输入...")
            
            while True:
                self.logger.debug("等待读取stdin输入...")
                line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
                self.logger.debug(f"读取到输入: {line[:100]}...")  # 只记录前100个字符
                
                if not line:
                    self.logger.info("stdin输入结束")
                    break
                    
                line = line.strip()
                if not line:
                    self.logger.debug("输入为空行，跳过处理")
                    continue
                
                try:
                    self.logger.debug("解析JSON消息")
                    message = json.loads(line)
                    self.logger.info(f"收到消息: {message.get('method', '未知方法')}")
                    await self.handle_request(message)
                    self.logger.debug("消息处理完成")
                        
                except json.JSONDecodeError as e:
                    error_msg = {
                        "jsonrpc": "2.0",
                        "error": {
                            "code": -32700,
                            "message": f"JSON解析错误: {str(e)}"
                        }
                    }
                    self.logger.error(f"JSON解析错误: {e}")
                    await self._send_json(error_msg)
                except Exception as e:
                    self.logger.error(f"处理消息异常: {e}", exc_info=True)
                    
        except KeyboardInterrupt:
            self.logger.info("训练工具正常退出")
        except Exception as e:
            self.logger.error(f"训练工具运行错误: {e}", exc_info=True)
        finally:
            self.logger.info("训练工具资源清理完成")


async def main():
    """MCP训练工具主函数"""
    # 初始化MCP工具
    mcp_tool = MiniLLMTrainingMCP()
    
    # 创建MCP服务器
    from mcp import Server
    server = Server("mini-llm-training-tool")
    
    # 注册工具
    @server_TOOL
    async def run_fine_tuning(params: Dict[str, Any]) -> Dict[str, Any]:
        return await mcp_tool.run_fine_tuning(params)
    
    @server_TOOL
    async def analyze_dataset(params: Dict[str, Any]) -> Dict[str, Any]:
        return await mcp_tool.analyze_dataset(params)
    
    @server_TOOL
    async def get_training_status(params: Dict[str, Any]) -> Dict[str, Any]:
        return await mcp_tool.get_training_status(params)
    
    # 运行MCP服务器
    await server.run()

# 导出可被调用的main函数
__all__ = ['main']

def run():
    """运行主函数"""
    asyncio.run(main())


if __name__ == "__main__":
    # 当脚本直接运行时启动工具
    run()